{-# LANGUAGE ScopedTypeVariables #-}
module Test.Snarkjs (tests, run) where

import Keelung
import Test.Hspec
import System.Process
import Data.Either (isRight)
import Data.Functor ((<&>))
import GHC.IO.Exception (ExitCode(ExitFailure))
import System.Directory (withCurrentDirectory, createDirectoryIfMissing)

run :: IO ()
run = hspec tests

tests :: SpecWith ()
tests = do
  let inDir = withCurrentDirectory "snarkjs"
  describe ("Testing Snarkjs-related compiler functions...\n[ INFO ] These tests are possible only if you have Keelung compiler >= " <> keelungVersion <> " and Snarkjs installed.") $ do
    _ <- runIO (createDirectoryIfMissing True "snarkjs")
    it "Generating .r1cs file (echo)" $ do
      shouldReturn (inDir (genCircuitBin "echo.r1cs" bn128 echo) <&> isRight) True
    it "Generating .r1cs file (quad)" $ do
      shouldReturn (inDir (genCircuitBin "quad.r1cs" bn128 quad) <&> isRight) True
    it "Generating .wtns file (echo)" $ do
      shouldReturn (inDir (genWtns "echo.wtns" bn128 echo [] [1]) <&> isRight) True
    it "Generating .wtns file (quad)" $ do
      shouldReturn (inDir (genWtns "quad.wtns" bn128 quad [3,5,-22] [2]) <&> isRight) True

    return ()
  describe "Testing compatibility with Snarkjs" $ do
    it "Checking if Snarkjs exists" $ do
      shouldReturn (readProcessWithExitCode "snarkjs" [] [] <&> (\(a,_,_) -> a)) (ExitFailure 99)
    it "Generating .ptau files with Snarkjs" $ do
      shouldReturn (setupSnarkjsPtau 10 "abcd" <&> isRight) True
    it "Generating .zkey with Snarkjs (echo)" $ do
      shouldReturn (genZkeyBySnarkjs PLONK "snarkjs/echo.r1cs" "snarkjs/pot10_final.ptau" <&> isRight) True
    it "Generating .zkey (quad)" $ do
      shouldReturn (genZkeyBySnarkjs PLONK "snarkjs/quad.r1cs" "snarkjs/pot10_final.ptau" <&> isRight) True
    it "Checking correctness of .wtns generated by Keelung with Snarkjs (echo)" $ do
      shouldReturn (checkWtnsBySnarkjs "snarkjs/echo.r1cs" "snarkjs/echo.wtns" <&> isRight) True
    it "Checking correctness of .wtns generated by Keelung with Snarkjs (quad)" $ do
      shouldReturn (checkWtnsBySnarkjs "snarkjs/quad.r1cs" "snarkjs/quad.wtns" <&> isRight) True
    it "Proving Keelung\'s .r1cs and .wtns by Snarkjs (echo)" $ do
      shouldReturn (proveBySnarkjs PLONK "snarkjs/echo.zkey" "snarkjs/echo.wtns" <&> isRight) True
    it "Proving Keelung\'s .r1cs and .wtns by Snarkjs (quad)" $ do
      shouldReturn (proveBySnarkjs PLONK "snarkjs/quad.zkey" "snarkjs/quad.wtns" <&> isRight) True

-- | A program that outputs whatever number is given.
echo :: Comp Field
echo = input Private -- request for a private input and return it

quad :: Comp ()
quad = do
    _a :: Field <- input Public
    _b :: Field <- input Public
    _c :: Field <- input Public
    _x :: Field <- input Private
    assert (eq (_a * _x * _x + _b * _x + _c) 0)
    return ()
